package com.duoec.dw.lineage.service.analyser.statement;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
import com.alibaba.druid.sql.ast.statement.SQLInsertInto;
import com.alibaba.druid.sql.ast.statement.SQLSelect;
import com.alibaba.druid.sql.ast.statement.SQLWithSubqueryClause;
import com.alibaba.druid.sql.dialect.hive.ast.HiveInsertStatement;
import com.duoec.dw.lineage.dto.*;
import com.duoec.dw.lineage.enums.StatementType;
import com.duoec.dw.lineage.service.LineageTableService;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.util.LinkedHashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * 表数据插入分析器
 *
 * @author xuwenzhen
 */
@Service
public class InsertStatementAnalyser extends BaseStatementAnalyser<HiveInsertStatement> {
    private static final Logger logger = LoggerFactory.getLogger(InsertStatementAnalyser.class);

    @Autowired
    private LineageTableService lineageTableService;

    /**
     * 匹配具体的SQLStatement
     *
     * @return 是否匹配
     */
    @Override
    public Set<Class<? extends SQLStatement>> parseSqlStatements() {
        return Sets.newHashSet(HiveInsertStatement.class);
    }

    /**
     * 分析
     *
     * @param statement SQLStatement
     * @return 当前insert语句的血缘关系
     */
    @Override
    public TableLineageDto analyse(HiveInsertStatement statement) {
        //with-as
        SQLWithSubqueryClause withSubqueryClause = statement.getWith();
        MultiTableLineageDto withAsTableLineages;
        if (withSubqueryClause != null) {
            withAsTableLineages = hiveAnalyser.analyseFrom(withSubqueryClause, null);
        } else {
            withAsTableLineages = new MultiTableLineageDto();
        }

        //query
        SQLSelect sqlSelect = statement.getQuery();
        TableLineageDto dto = hiveAnalyser.analyseQuery(sqlSelect.getQuery(), withAsTableLineages);

        //获取表坐标
        TableCoordinateDto tableCoordinate = getTableCoordinate(statement);
        dto.setDatabaseName(tableCoordinate.getDatabaseName());
        dto.setTableName(tableCoordinate.getTableName());

        //加载目标表信息
        TableInfoDto tableInfoDto = lineageTableService.getTableDto(tableCoordinate);
        if (tableInfoDto == null) {
            return null;
        }
        dto.setComment(tableInfoDto.getComment());

        //partitions
        hiveAnalyser.analysePartitions(dto, statement.getPartitions());

        LinkedHashSet<TableFieldDto> fields = tableInfoDto.getFields();
        List<TableFieldLineageDto> lineageFields = Lists.newArrayList(dto.getFields());

        int lineageFieldCount = getLineageFieldCount(dto);
        if (fields.size() > lineageFieldCount) {
            logger.warn("表[{}]schema字段数量({})与血缘关系分析出来的字段数量({})不匹配！", tableCoordinate, fields.size(), lineageFieldCount);
        } else if (fields.size() < lineageFieldCount) {
            logger.error("表[{}]schema字段数量({})小于血缘关系分析出来的字段数量({})！", tableCoordinate, fields.size(), lineageFieldCount);
        }

        resetLineageFieldNames(fields, lineageFields);

        dto.setStatementType(StatementType.INSERT_TABLE_DATA);
        return dto;
    }

    public int getLineageFieldCount(TableLineageDto dto) {
        Set<String> fields;
        if (!CollectionUtils.isEmpty(dto.getFields())) {
            fields = dto.getFields().stream().map(TableFieldDto::getColName).collect(Collectors.toSet());
        } else {
            fields = Sets.newHashSet();
        }
        if (!CollectionUtils.isEmpty(dto.getPartitions())) {
            fields.addAll(dto.getPartitions());
        }
        return fields.size();
    }

    private void resetLineageFieldNames(LinkedHashSet<TableFieldDto> fields, List<TableFieldLineageDto> lineageFields) {
        int i = 0;
        for (TableFieldDto tableField : fields) {
            String colName = tableField.getColName();
            if (lineageFields.size() > i) {
                TableFieldLineageDto fieldLineage = lineageFields.get(i);
                String lineageFieldName = fieldLineage.getColName();
                if (!colName.equals(lineageFieldName)) {
                    if (lineageFieldName == null || lineageFieldName.contains("(")) {
                        //直接替换
                        fieldLineage.setColName(colName);
                    } else {
                        logger.warn("[{}]字段不一致：{}:{}", i, colName, lineageFieldName);
                    }
                }
            }
            ++i;
        }
    }

    /**
     * 获取表坐标
     *
     * @param statement SQLInsertInto
     * @return 表坐标
     */
    private TableCoordinateDto getTableCoordinate(SQLInsertInto statement) {
        TableCoordinateDto tableCoordinate = new TableCoordinateDto();

        SQLExprTableSource tableSource = statement.getTableSource();
        SQLExpr expr = tableSource.getExpr();
        if (expr instanceof SQLPropertyExpr) {
            SQLPropertyExpr sqlPropertyExpr = (SQLPropertyExpr) expr;
            tableCoordinate.setDatabaseName(sqlPropertyExpr.getOwnernName());
            tableCoordinate.setTableName(sqlPropertyExpr.getName());
        }

        return tableCoordinate;
    }

    public void setLineageTableService(LineageTableService lineageTableService) {
        this.lineageTableService = lineageTableService;
    }
}
